from huggingface_hub import HfApi
from concurrent.futures import ThreadPoolExecutor
import os
import time

# 设置常量，避免硬编码
HUB_API_ENDPOINT = "https://huggingface.sukaka.top"

script_dir = os.path.dirname(os.path.realpath(__file__))

def download_hf_and_replace_links(repo, folder):
    tasks = []
    hf_api = HfApi(endpoint=HUB_API_ENDPOINT)

    retry_count = 3
    for _ in range(retry_count):
        try:
            files = hf_api.list_repo_files(repo)
            files = [fi for fi in files if fi not in ['.gitattributes', 'README.md']]
            break
        except Exception as e:
            print(f"Error fetching files for {repo}: {e}")
            time.sleep(5)
    else:
        print(f"Failed to fetch files for {repo} after {retry_count} retries. Exiting.")
        return

    for fi in files:
        li = os.path.join(HUB_API_ENDPOINT, repo, "resolve/main/", fi)
        tasks.append(f'aria2c --console-log-level=error -c -x 16 -s 16 -k 1M -t 10 -d {folder} -o {fi} "{li}" > /dev/null 2>&1')
    
    return tasks

def execute_and_output(command):
    print(f"Executing: {command}")
    os.system(command)

def download_hf_repo_files(repo_list, base_path):
    with ThreadPoolExecutor() as executor:
        for repo in repo_list:
            # 获取仓库名称的最后一部分作为目录名
            repo_name = repo.split("/")[-1]
            folder = os.path.join(base_path, repo_name)
            os.makedirs(folder, exist_ok=True)
            
            tasks = download_hf_and_replace_links(repo, folder)
            
            for task in tasks:
                executor.submit(execute_and_output, task)

# 示例用法，repo_list 包含多个仓库
repo_list = ["yzd-v/DWPose"]
# 构建基础路径
base_path = os.path.join(script_dir, "video_controlnet_aux", "ckpts", "yzd-v", "DWPose")
download_hf_repo_files(repo_list, base_path)